import serial
import rclpy
from sensor_msgs.msg import Imu
# from tkinter import *
import struct
import codecs
from datetime import datetime
import csv
from rclpy.node import Node
class IMUInfo:  # 接收数据初始化
    def __init__(self):
        # 加速度
        self.accel_x = 0.0
        self.accel_y = 0.0
        self.accel_z = 0.0
        # 角速度
        self.angle_x = 0.0
        self.angle_y = 0.0
        self.angle_z = 0.0
        # 欧拉角
        self.pitch = 0.0
        self.roll = 0.0
        self.yaw = 0.0
        # 四元数
        self.quaternion_data0 = 0.0
        self.quaternion_data1 = 0.0
        self.quaternion_data2 = 0.0
        self.quaternion_data3 = 0.0
        # 温度
        self.temperature = 0.0
        # 采样时间戳
        self.sample_timestamp = 0
        # 同步输出时间戳
        self.data_ready_timestamp = 0

class Ros2IMUDriver(Node):
    def __init__(self):
        super().__init__('imu_driver')
        
        # wangshuo, 20250315 add Imu pub
        self.imu_pub = self.create_publisher(Imu, 'imu_board', 10)

        self.frame_id = self.get_frame_id()
    def Read_IMUdata(self,data):

        # 定义解析头
        PROTOCOL_FIRST_BYTE = 0x59
        PROTOCOL_SECOND_BYTE = 0x53
        # 定义数据因子
        fac1 = 0.000001
        fac2 = 0.001

        if data[0]==0x59 and data[1]==0x53:
            imu_data = IMUInfo()
            print('\n**********开始解析数据**********')
            print('解析头校验：',data[0] == PROTOCOL_FIRST_BYTE and data[1] == PROTOCOL_SECOND_BYTE)
            # 2.分析数据域
            # 输出顺序为    IMU温度  加速度  角速度  磁场归一化值 磁场强度 欧拉角 四元数 采样时间戳 同步输出时间戳
            # 数据类型为      2+2    2+12   2+12     2+12      2+12  2+12  2+16    2+4        2+4
            # payload_len=40 (14+14+6+6)(加速度 角速度 采样时间戳 同步输出时间戳)
            header = struct.unpack('HHB', data[:5]) # 解包协议头部分：帧头、帧序号、数据长度
            payload_len = header[2] # 返回数据长度
            print('数据长度：',payload_len)
            if payload_len>0:
                data_ = data[5:]
                print('各数据id及长度',data_[:2],data_[14:16],data_[28:30],data_[42:44],data_[60:62],data_[64:66],data_[70:72])
                # 解析加速度(0-13)
                if data_[0]==0x10 and data_[1]==12:
                    # print("解析加速度")
                    # 每四个字节组成一个int型（有符号）数据，小端模式(最低有效字节排在最前面)，然后乘以系数
                    imu_data.accel_x = int.from_bytes(data_[2:6],byteorder='little',signed = True)*fac1
                    imu_data.accel_y = int.from_bytes(data_[6:10],byteorder='little',signed = True)*fac1
                    imu_data.accel_z = int.from_bytes(data_[10:14],byteorder='little',signed = True)*fac1

                # 解析角速度(14-27)
                if data_[14]==0x20 and data_[15]==12:

                    # print("解析角速度")
                    imu_data.angle_x = int.from_bytes(data_[16:20],byteorder='little',signed = True)*fac1
                    imu_data.angle_y = int.from_bytes(data_[20:24],byteorder='little',signed = True)*fac1
                    imu_data.angle_z = int.from_bytes(data_[24:28],byteorder='little',signed = True)*fac1
                    
                # 解析欧拉角(28-41)
                if data_[28]==0x40 and data_[29]==12:
                    # print("解析欧拉角")
                    imu_data.pitch = int.from_bytes(data_[30:34],byteorder='little',signed = True)*fac1
                    imu_data.roll = int.from_bytes(data_[34:38],byteorder='little',signed = True)*fac1
                    imu_data.yaw = int.from_bytes(data_[38:42],byteorder='little',signed = True)*fac1

                # 解析四元数(42-59)
                if data_[42]==0x41 and data_[43]==16:

                    # print("解析四元数")
                    imu_data.quaternion_data0 = int.from_bytes(data_[44:48],byteorder='little',signed = True)*fac1
                    imu_data.quaternion_data1 = int.from_bytes(data_[48:52],byteorder='little',signed = True)*fac1
                    imu_data.quaternion_data2 = int.from_bytes(data_[52:56],byteorder='little',signed = True)*fac1
                    imu_data.quaternion_data3 = int.from_bytes(data_[56:60],byteorder='little',signed = True)*fac1
                
                # 解析温度(60-65)
                if data_[60]==0x01 and data_[61]==2:
                    # print("解析温度")
                    imu_data.temperature = int.from_bytes(data_[62:64],byteorder='little',signed = True)*fac2

                # 解析采样时间戳(64-69)
                if data_[64]==0x51 and data_[65]==4:
                    # print("解析采样时间戳")
                    imu_data.sample_timestamp = int.from_bytes(data_[66:70],byteorder='little',signed = True)

                # 解析采样时间戳(70-75)
                if data_[70]==0x52 and data_[71]==4:
                    # print("解析同步输出时间戳")
                    imu_data.data_ready_timestamp = int.from_bytes(data_[72:],byteorder='little',signed = True)

                print('加速度(m/s²)：',imu_data.accel_x,imu_data.accel_y,imu_data.accel_z)
                print('角速度(deg/s)：',imu_data.angle_x,imu_data.angle_y,imu_data.angle_z)
                print('欧拉角(deg)：',imu_data.pitch,imu_data.roll,imu_data.yaw)
                print('温度(°C)：',imu_data.temperature)
                print('四元数：',imu_data.quaternion_data0,imu_data.quaternion_data1,imu_data.quaternion_data2,imu_data.quaternion_data3)
                print('采样时间戳(us)：',imu_data.sample_timestamp)
                print('同步输出时间戳(us)：', imu_data.data_ready_timestamp)

                # 获取当前时间
                now = datetime.now()
                # 格式化为包含毫秒的字符串
                formatted_time = now.strftime("%Y.%m.%d.%H.%M.%S.%f")[:-3]
                print('系统时间戳：',formatted_time)

                imu_msg = Imu()
                imu_msg.header.stamp = self.get_clock().now().to_msg()
                imu_msg.header.frame_id = self.frame_id
                imu_msg.linear_acceleration.x = imu_data.accel_x
                imu_msg.linear_acceleration.y = imu_data.accel_y
                imu_msg.linear_acceleration.z = imu_data.accel_z

                imu_msg.angular_velocity.x = imu_data.angle_x
                imu_msg.angular_velocity.y = imu_data.angle_y
                imu_msg.angular_velocity.z = imu_data.angle_z

                imu_msg.orientation.x = imu_data.quaternion_data0
                imu_msg.orientation.y = imu_data.quaternion_data1
                imu_msg.orientation.z = imu_data.quaternion_data2
                imu_msg.orientation.w = imu_data.quaternion_data3
                self.imu_pub.publish(imu_msg)
                print("publish imu success!")

        # 当协议头判断失败，则重新打开串口
            return True
        else:
            self.get_logger().warn("协议头判断失败, 重新打开串口")
            return False
    def get_frame_id(self):
        frame_id = self.declare_parameter('frame_id', 'gps').value
        prefix = self.declare_parameter('tf_prefix', '').value
        if len(prefix):
            return '%s/%s' % (prefix, frame_id)
        return frame_id
